package com.zzh.zzhaiagent.memory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.serializers.CollectionSerializer;
import com.esotericsoftware.kryo.util.Pool;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class FileBasedChatMemory implements ChatMemory {

    private final Map<String, List<Message>> conversationHistory = new ConcurrentHashMap<>();
    private final Path storageDirectory;
    private final Pool<Kryo> kryoPool;

    public FileBasedChatMemory(Path storageDirectory) {
        this.storageDirectory = storageDirectory;
        this.kryoPool = createKryoPool();

        try {
            Files.createDirectories(storageDirectory); // 确保目录存在
        } catch (IOException e) {
            throw new RuntimeException("Failed to create storage directory", e);
        }
    }

    private Pool<Kryo> createKryoPool() {
        return new Pool<>(true, false) {
            @Override
            protected Kryo create() {
                Kryo kryo = new Kryo();
                kryo.setRegistrationRequired(false);
                kryo.register(ArrayList.class, new CollectionSerializer());
                kryo.register(Message.class); // 假设Message类已实现Kryo序列化
                return kryo;
            }
        };
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        conversationHistory.compute(conversationId, (k, v) -> {
            List<Message> list = (v == null) ? new ArrayList<>() : v;
            list.addAll(messages);
            return list;
        });
        saveConversation(conversationId);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> messages = conversationHistory.computeIfAbsent(conversationId, this::loadConversation);
        int fromIndex = Math.max(0, messages.size() - lastN);
        return new ArrayList<>(messages.subList(fromIndex, messages.size()));
    }

    @Override
    public void clear(String conversationId) {
        conversationHistory.remove(conversationId);
        try {
            Files.deleteIfExists(resolveFilePath(conversationId));
        } catch (IOException e) {
            throw new RuntimeException("Failed to clear conversation: " + conversationId, e);
        }
    }

    // 核心存储逻辑
    private void saveConversation(String conversationId) {
        List<Message> messages = conversationHistory.get(conversationId);
        if (messages == null) return;

        Kryo kryo = kryoPool.obtain();
        try (Output output = new Output(new FileOutputStream(resolveFilePath(conversationId).toFile()))) {
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            throw new RuntimeException("Failed to save conversation: " + conversationId, e);
        } finally {
            kryoPool.free(kryo);
        }
    }

    private List<Message> loadConversation(String conversationId) {
        Path filePath = resolveFilePath(conversationId);
        if (!Files.exists(filePath)) return new ArrayList<>();

        Kryo kryo = kryoPool.obtain();
        try (Input input = new Input(new FileInputStream(filePath.toFile()))) {
            return kryo.readObject(input, ArrayList.class);
        } catch (IOException e) {
            throw new RuntimeException("Failed to load conversation: " + conversationId, e);
        } finally {
            kryoPool.free(kryo);
        }
    }

    // 文件名安全处理
    private Path resolveFilePath(String conversationId) {
        String safeFileName = Base64.getUrlEncoder().encodeToString(
            conversationId.getBytes(StandardCharsets.UTF_8)
        ) + ".bin";
        return storageDirectory.resolve(safeFileName);
    }
}